# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
from copy import deepcopy
from pathlib import Path

import pytest

from olive.engine.footprint import Footprint

# pylint: disable=attribute-defined-outside-init, protected-access


class TestFootprint:
    @pytest.fixture(autouse=True)
    def setup(self):
        self.footprint_file = Path(__file__).parent / "mock_data" / "footprints.json"
        self.fp = Footprint.from_file(self.footprint_file)
        self.input_node = {k: v for k, v in self.fp.nodes.items() if v.parent_model_id is None}

    def test_create_from_model_ids(self):
        self.fp.set_output_model_ids(list(self.fp.nodes.keys()))
        assert len(self.fp.output_model_ids) == len(self.fp.nodes)
        assert self.fp.output_model_ids == list(self.fp.nodes.keys())

    def test_file_dump(self, tmp_path):
        self.fp.to_file(tmp_path / "footprint.json")
        fp2 = Footprint.from_file(tmp_path / "footprint.json")
        assert len(fp2.nodes) == 3

    def test_json_dump(self):
        json_fp = self.fp.to_json()
        fp2 = Footprint.from_json(json_fp)
        assert len(fp2.nodes) == 3

    def test_pareto_frontier(self):
        self.fp.create_pareto_frontier()
        assert len(self.fp.output_model_ids) == 2
        assert all(self.fp.nodes[model_id].is_pareto_frontier for model_id in self.fp.output_model_ids)

    def test_empty_pareto_frontier(self):
        # set all `is_goals_met` as false
        unmet_goals_nodes = deepcopy(self.fp.nodes)
        for v in unmet_goals_nodes.values():
            v.metrics.if_goals_met = False
        new_fp = Footprint(nodes=unmet_goals_nodes)
        new_fp.create_pareto_frontier()
        assert len(new_fp.output_model_ids) == 0

    def test_trace_back_run_history(self):
        for model_id in self.fp.nodes:
            run_history = self.fp.trace_back_run_history(model_id)
            assert run_history is not None
            assert next(reversed(run_history)) in self.input_node

    def test_get_model_info(self):
        for model_id in self.fp.nodes:
            inference_config = self.fp.get_model_inference_config(model_id)
            model_path = self.fp.get_model_path(model_id)
            if model_id in self.input_node:
                assert inference_config is None
                assert str(model_path).endswith(".pt")
            else:
                assert inference_config is not None
                assert str(model_path).endswith(".onnx")

    def test_plot_pareto_frontier(self, tmp_path):
        self.fp.objective_dict = {
            "accuracy-accuracy_score": {"higher_is_better": True, "priority": 1},
            "latency-avg": {"higher_is_better": False, "priority": 2},
        }
        self.fp._plot_pareto_frontier(
            is_show=False,
            save_path=tmp_path / "pareto_frontier.html",
        )
        assert (tmp_path / "pareto_frontier.html").exists()
